package tree;

public class RedBlackTree<Key extends Comparable<Key>, Value> {
    // 根节点
    private Node root;
    // 个数
    private int N;
    //红色链接
    private static final boolean RED = true;
    //黑色链接
    private static final boolean BLACK = false;


    private boolean isRed(Node x) {
        //空结点默认是黑色链接
        if (x == null) {
            return false;
        }
        return x.color == RED;
    }

    /**
     * 左旋
     */
    private Node rotateLeft(Node h) {
        // 当前h 的右
        Node hRight = h.right;
        // 当前h 的右的左
        Node lhRight = hRight.left;
        // 让 h的右 指向 h的右的左
        h.right = lhRight;
        // 让 原来h右 的左指向当前h
        hRight.left = h;
        // 原来h右 变成 当前h的颜色
        hRight.color = h.color;
        // 让h 变红
        h.color = RED;
        // 返回h右
        return hRight;
    }

    /**
     * 右旋
     */
    private Node rotateRight(Node h) {
        // 当前h的左
        Node hLeft = h.left;
        // 当前h的左的右
        Node rhLeft = hLeft.right;
        // 让当前h的左的右变成h的左
        h.left = rhLeft;
        // 让当前h 变成当前左的右
        hLeft.right = h;
        // 颜色反转
        hLeft.color = h.color;
        h.color = RED;
        return hLeft;
    }

    /**
     * 颜色反转，相当于完成4-结点的拆分
     */
    private void flipColors(Node h) {
        // 自己变成红
        h.color = RED;
        // 左右变黑
        h.left.color = BLACK;
        h.right.color = BLACK;
    }


    /**
     * 插入
     */
    public void put(Key key, Value value) {
        //在root整个树上插入key-val
        root = put(root, key, value);
        // 每次插入都让根结点的颜色变为BLACK
        root.color = BLACK;
    }

    private Node put(Node h, Key key, Value value) {
        if (h == null) {
            // 标准的插入操作 和 以红链连接父节点
            N++;
            return new Node(key, value, null, null, RED);
        }
        int cmp = key.compareTo(h.key);
        if (cmp > 0) {
            // 大于 往右走
            h.right = put(h.right, key, value);
        } else if (cmp < 0) {
            // 小于 往左走
            h.left = put(h.left, key, value);
        } else {
            //等于 直接赋值
            h.value = value;
        }

        // 开始旋转已经颜色反转
        // 1. 左黑右红 左旋
        if (!isRed(h.left) && isRed(h.right)) {
            h = rotateLeft(h);
        }
        // 2. 左红 左左红 右旋
        if (isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }
        // 3. 左右红 反转
        if (isRed(h.left) && isRed(h.right)) {
            flipColors(h);
        }

        return h;
    }


    public Value get(Key key) {
        return get(root, key);
    }

    private Value get(Node root, Key key) {
        Node n = root;
        int cmp = key.compareTo(n.key);
        if (cmp == 0) {
            return n.value;
        } else if (cmp > 0) {
            return get(root.right, key);
        } else {
            return get(root.left, key);
        }
    }

    //获取树中元素的个数
    public int size() {
        return N;
    }


    //结点类
    private class Node {
        //存储键
        public Key key;
        //存储值
        public Value value;
        //记录左子结点
        public Node left;
        //记录右子结点
        public Node right;
        //由其父结点指向它的链接的颜色
        public boolean color;

        public Node(Key key, Value value, Node left, Node right, boolean color) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            this.color = color;
        }
    }

    public static void main(String[] args) {
        RedBlackTree<Integer, String> bt = new RedBlackTree<>();
        bt.put(4, "二哈");
        bt.put(1, "张三");
        bt.put(3, "李四");
        bt.put(5, "王五");
        System.out.println(bt.size());
        bt.put(1, "老三");
        System.out.println(bt.get(1));
        System.out.println(bt.size());
    }
}
